# -*- coding: utf-8 -*-
"""
企业级知识库问答系统
"""
import os
import time
import uuid
from typing import List, Dict, Optional
import dashscope
from dashscope import Generation
from sentence_transformers import SentenceTransformer

from config import MODEL_CONFIG, SYSTEM_CONFIG, RAG_CONFIG
from utils.logger import setup_logger
from utils.doc_processor import DocumentProcessor
from utils.vector_store import VectorStore
from utils.session import SessionManager

logger = setup_logger(__name__)

class QAAssistant:
    def __init__(self, api_key: str = None):
        """
        初始化问答助手
        
        Args:
            api_key: 通义千问API密钥
        """
        # 设置API密钥
        if api_key:
            dashscope.api_key = api_key
        else:
            dashscope.api_key = os.getenv('DASHSCOPE_API_KEY')
            
        if not dashscope.api_key:
            raise ValueError("请设置通义千问API密钥")
            
        logger.info("初始化问答助手...")
        
        # 初始化组件
        self.doc_processor = DocumentProcessor()
        self.vector_store = VectorStore()
        self.session_manager = SessionManager()
        self.embedding_model = SentenceTransformer(MODEL_CONFIG['embedding_model'])
        
        # 设置提示词模板
        self.prompt_template = """
        你是一个专业的企业知识库问答助手。我会给你一些相关的文档片段、对话历史和一个问题。
        请你基于这些信息，以专业、准确、简洁的方式回答问题。
        如果文档片段中没有足够的信息来回答问题，请明确指出。
        
        相关文档片段：
        {context}
        
        最近的对话历史：
        {history}
        
        当前问题：{question}
        
        请给出你的回答：
        """
        
        logger.info("问答助手初始化完成")
    
    def process_documents(self, dir_path: str = "docs"):
        """处理文档并存入向量库"""
        try:
            logger.info("开始处理文档...")
            
            # 处理文档
            documents = self.doc_processor.process_directory(dir_path)
            
            if not documents:
                logger.warning("没有找到需要处理的文档")
                return
            
            # 生成嵌入向量
            texts = [doc['text'] for doc in documents]
            embeddings = self.embedding_model.encode(texts, show_progress_bar=True)
            
            # 清空现有向量库并添加新文档
            self.vector_store.clear()
            self.vector_store.add_documents(documents, embeddings.tolist())
            
            logger.info(f"文档处理完成，共处理 {len(documents)} 个文本片段")
            
        except Exception as e:
            logger.error(f"处理文档时出错: {str(e)}")
            raise
    
    def _get_answer_with_retry(self, prompt: str) -> str:
        """带重试机制的API调用"""
        max_retries = SYSTEM_CONFIG['max_retries']
        retry_delay = SYSTEM_CONFIG['retry_delay']
        
        for attempt in range(max_retries):
            try:
                response = Generation.call(
                    model=MODEL_CONFIG['llm_model'],
                    prompt=prompt,
                    temperature=MODEL_CONFIG['temperature'],
                    top_p=MODEL_CONFIG['top_p'],
                    result_format='message',
                    max_tokens=MODEL_CONFIG['max_tokens']
                )
                
                if response.status_code == 200:
                    if (hasattr(response, 'output') and 
                        hasattr(response.output, 'choices') and 
                        response.output.choices and 
                        hasattr(response.output.choices[0], 'message') and
                        hasattr(response.output.choices[0].message, 'content')):
                        
                        return response.output.choices[0].message.content
                    else:
                        logger.error("API响应格式错误")
                        raise ValueError("API响应格式错误")
                else:
                    error_msg = getattr(response, 'message', '未知错误')
                    logger.error(f"API调用失败: {error_msg}")
                    raise ValueError(f"API调用失败: {error_msg}")
                    
            except Exception as e:
                if attempt < max_retries - 1:
                    logger.warning(f"第{attempt + 1}次调用失败，{retry_delay}秒后重试: {str(e)}")
                    time.sleep(retry_delay)
                else:
                    raise
    
    def get_answer(self, question: str, session_id: str = None) -> str:
        """获取问题的答案"""
        try:
            # 确保有会话ID
            session_id = session_id or str(uuid.uuid4())
            session = self.session_manager.get_session(session_id)
            
            # 获取问题的嵌入向量
            query_embedding = self.embedding_model.encode(question)
            
            # 搜索相关文档 - 增加返回数量
            similar_docs = self.vector_store.search_similar(
                query_embedding.tolist(),
                top_k=RAG_CONFIG['max_docs_per_query'],
                min_score=RAG_CONFIG['min_similarity_score']
            )
            
            # 构建提示词
            if similar_docs:
                # 获取扩展上下文
                context = self._get_extended_context(similar_docs)
            else:
                context = "未找到相关文档。"
            
            # 获取历史对话
            history = session.get_recent_context()
            
            # 构建完整提示词
            prompt = self.prompt_template.format(
                context=context,
                history=history,
                question=question
            )
            
            # 获取答案
            answer = self._get_answer_with_retry(prompt)
            
            # 记录交互
            session.add_interaction(question, answer, similar_docs)
            
            return answer
            
        except Exception as e:
            logger.error(f"获取答案时出错: {str(e)}")
            return f"抱歉，处理您的问题时出现错误: {str(e)}"
    
    def _get_extended_context(self, similar_docs: List[Dict]) -> str:
        """获取扩展的上下文内容"""
        # 按文档和页码分组
        docs_by_source = {}
        for doc in similar_docs:
            key = (doc['metadata']['source'], doc['metadata']['page'])
            if key not in docs_by_source:
                docs_by_source[key] = []
            docs_by_source[key].append(doc)
        
        # 构建扩展上下文
        context_parts = []
        total_length = 0
        
        for (source, page), docs in docs_by_source.items():
            # 获取当前页的所有段落
            page_segments = self.vector_store.get_page_segments(source, page)
            
            # 找到相关段落在页面中的位置
            relevant_indices = [page_segments.index(doc['text']) for doc in docs if doc['text'] in page_segments]
            
            # 扩展上下文窗口
            extended_indices = set()
            for idx in relevant_indices:
                start = max(0, idx - RAG_CONFIG['context_window'])
                end = min(len(page_segments), idx + RAG_CONFIG['context_window'] + 1)
                extended_indices.update(range(start, end))
            
            # 构建上下文
            page_context = []
            for idx in sorted(extended_indices):
                text = page_segments[idx]
                # 检查是否超过最大长度限制
                if total_length + len(text) > RAG_CONFIG['max_context_length']:
                    break
                page_context.append(text)
                total_length += len(text)
            
            if page_context:
                context_parts.append(
                    f"[来源: {source} 第{page}页]\n" +
                    "\n".join(page_context)
                )
        
        return "\n\n".join(context_parts)

def chat_loop():
    """命令行交互循环"""
    try:
        # 初始化助手
        assistant = QAAssistant()
        session_id = str(uuid.uuid4())
        
        print("\n=== 欢迎使用企业知识库问答系统 ===")
        print("命令说明:")
        print("1. 输入 'quit' 或 'exit' 退出对话")
        print("2. 输入 'clear' 清空屏幕")
        print("3. 输入 'reload' 重新加载知识库")
        print("4. 输入 'help' 显示帮助信息")
        print("支持的文档格式: PDF, TXT, DOC, DOCX, MD")
        print("请将文档放在 docs 目录中")
        print("====================================\n")
        
        # 首次加载文档
        print("正在加载知识库...")
        assistant.process_documents("docs")
        print("知识库加载完成！\n")
        
        while True:
            # 获取用户输入
            question = input("\n请输入您的问题: ").strip()
            
            # 处理特殊命令
            if question.lower() in ['quit', 'exit', 'q']:
                print("\n感谢使用，再见！")
                break
            
            if question.lower() == 'clear':
                os.system('clear' if os.name == 'posix' else 'cls')
                continue
            
            if question.lower() == 'reload':
                print("\n正在重新加载知识库...")
                assistant.process_documents()
                print("知识库重新加载完成！")
                continue
            
            if question.lower() == 'help':
                print("\n=== 帮助信息 ===")
                print("1. 输入问题即可获得答案")
                print("2. 系统会自动在知识库中搜索相关内容")
                print("3. 如果没有找到相关内容，会基于模型的通用知识回答")
                print("4. 使用 'reload' 命令可以重新加载最新的文档")
                print("===============\n")
                continue
            
            # 空输入则继续
            if not question:
                continue
            
            # 获取答案
            print("\n正在思考...")
            answer = assistant.get_answer(question, session_id)
            print("\n回答：", answer)
            
    except KeyboardInterrupt:
        print("\n\n程序被中断，正在退出...")
    except Exception as e:
        logger.error(f"程序运行出错: {str(e)}")
        print(f"\n发生错误: {str(e)}")

if __name__ == "__main__":
    chat_loop()